import util.mydecorator
import util.lib as lib
import util.util as utils

import pandas as pd

gb_dict = {
    "20w_gb": None,
    "30w_gb": None,
    "60w_gb": None,
}

gb_keys = ["20w_gb","30w_gb","60w_gb"]

# 计算差距
def cal_duration(t1,t2):
    return 52 * (int(t1[0:4]) - int(t2[0:4])) + int(t1[6:8]) - int(t2[6:8])

def cal_gb(df, key):

    global gb_dict

    df["next_week_pct"] = df["week_pct"].shift(-1)
    df["60dis_pct"] = df["dis_pct"]
    tmp = df[
        (df["close"] >= df[key]) & (df["close"].shift() < df[key].shift()) |
        (df["close"] < df[key]) & (df["close"].shift() >= df[key].shift())
        ].copy()
    if len(tmp) == 0:
        return

    tmp["type"] = key

    tmp["disPct"] = round(100 * (df["close"]/df[key] -1), 2)

    tmp["duration"] = 0
    tmp["revenge"] = 0.0
    tmp["revenge_pct"] = 0.0
    for i in range(len(tmp)):
        if i > 0:
            tmp.loc[tmp.iloc[i].name,"duration"] = (cal_duration(tmp.iloc[i].name, tmp.iloc[i-1].name)) * (1 if tmp.iloc[i]["disPct"] >= 0 else -1)
            tmp.loc[tmp.iloc[i].name,"revenge"] = (tmp.iloc[i].close - tmp.iloc[i-1].close) * (0 if tmp.iloc[i]["disPct"] >= 0 else 1)
            tmp.loc[tmp.iloc[i].name,"revenge_pct"] = round(100 *(tmp.iloc[i].close / tmp.iloc[i-1].close -1), 2) * (0 if tmp.iloc[i]["disPct"] >= 0 else 1)
        else :
            tmp.loc[tmp.iloc[i].name,"duration"] = len(df[:tmp.iloc[i].name]) * (1 if tmp.iloc[i]["disPct"] >= 0 else -1)
    
    if gb_dict[key] is None:
        gb_dict[key] = tmp
    else :
        #gb_dict[key] = gb_dict[key].append(tmp)
        gb_dict[key] = utils.merge(gb_dict[key],tmp, False)

def cal(k, df):
    for v in gb_keys:
        cal_gb(df, v)

@util.mydecorator.calTime
def job():

    global gb_dict

    writer = pd.ExcelWriter(lib.data_path + "gb_sta.xlsx")

    week_data = lib.lib_get_week_data()

    for k,v in week_data.items():
        cal(k, v)

    for v in gb_keys:
        gb_dict[v][["stock","category","wg_industry","close",
                           "15w_gb","20w_gb", "30w_gb","60w_gb",
                           "15dis_pct","20dis_pct", "30dis_pct","60dis_pct",
                           "level_point","rel_lp",
                           "week_pct" , "next_week_pct",
                           "disPct","duration", 
                           "revenge","revenge_pct","type"]].to_excel(writer, sheet_name=v)

    writer.close()


if __name__ == '__main__':
    job()